from flask import Flask, request, jsonify
from flask_cors import *
from flask_sslify import SSLify
from flask_socketio import SocketIO, emit

import random
import torch
import json
import math
import os

import a2c_FoV as a2c
import Arguments
from prepareFoVdata import *
from torch.distributions import Categorical

app = Flask(__name__)
app.config['DEBUG'] = True
app.config['SECRET_KEY'] = 'secret_key'
# sslify = SSLify(app)
# CORS(app, supports_credentials=True)
server_cert = './certificate/server/server-cert.pem'
server_key = './certificate/server/server-key.pem'

socketio = SocketIO()
socketio.init_app(app, cors_allowed_origins='*', async_mode='threading')

# ABR paremeters
S_INFO = 6
# bit_rate, buffer_size, rebuffering_time, bandwidth_measurement, chunk_til_video_end
S_LEN = 8  # take how many frames in the past
A_DIM = 4  #
ACTOR_LR_RATE = 0.0001
BUFFER_NORM_FACTOR = 10.0  # unit：s
CHUNK_TIL_VIDEO_END_CAP = 200
RAND_RANGE = 1000
M_IN_K = 1000.0
DEVIATION_THRESHOLD = 0.25
NN_MODEL = './model/actor_nn_model_ep_640.pkl'
actor = a2c.ActorNet(s_dim=[S_INFO, S_LEN], a_dim=A_DIM, lr=ACTOR_LR_RATE)
movie_json = None

# FoV parameters
videoName = "tiledSkiing"
videoId = 0
isLoad = [False for _ in range(len(videoList))]
sal_maps = None
time_array = None
FoV_net = convlstm.ConvLSTM_model(input_dim=args.input_size,
                                  hidden_dim=args.hidden_size,
                                  kernel_size=(5, 5),
                                  num_layers=args.num_layers,
                                  batch_first=True)

def convert_torch(variable, dtype=np.float32):
    if variable.dtype != dtype:
        variable = variable.astype(dtype)
    return torch.from_numpy(variable)

@socketio.on('connect', namespace='/ABR')
def connectABR():
    global movie_json
    # np.random.seed(20)
    actor.load_state_dict(torch.load(NN_MODEL))
    # 加载电影json数据
    with open('./model/tiledSkiing.json', 'r') as load_f:
        movie_json = json.load(load_f)
    socketio.emit('model response', {'data': "Success! Have Loaded ABR model parameters!"},
              namespace='/ABR')
    print("Success! Have Loaded ABR model parameters!")

@socketio.on('ABR request', namespace='/ABR')
# @app.route('/ABR', methods=['POST'])
def ABR(msg):
    face_idx = 0
    tile_idx = 0
    try:
        coord = msg['data']['coord']
    except TypeError:
        print("Get data some error osh")
        socketio.emit('ABR response', {'data': '0'}, to=msg['id'])
        return 0

    buffer = 0 if msg['data']['buffer'] is None else  msg['data']['buffer']
    lastQuality = msg['data']['lastQuality']
    past_8_index = msg['data']['past_8_index']
    past_8_throughput = msg['data']['past_8_throughput']
    tile_dev = msg['data']['tile_dev']
    movie_name = 'tiledSkiing'
    index = msg['data']['index']
    face = msg['data']['face']
    bitrates = movie_json['bitrate_bps']

    if movie_name == 'tiledSkiing':
        face_idx = face // 4
        tile_idx = face % 4

    fov_tmp = 1 if tile_dev <= 20 else 0

    past_download_time = np.zeros(8)
    for i in range(8):
        if past_8_throughput[i] == 0 or past_8_throughput[i] is None:
            past_8_throughput[i] = 0
            past_download_time[i] = 0
        else:
            if movie_name == 'tiledSkiing':
                past_download_time[i] = movie_json['segment_size_bytes'][max(0, index - 8 + i)][face_idx][tile_idx][
                                            past_8_index[i]] / past_8_throughput[i]
            else:
                past_download_time[i] = movie_json['segment_size_bytes'][max(0, index - 8 + i)][face][past_8_index[i]] / \
                                        past_8_throughput[i]
            past_8_throughput[i] /= M_IN_K
            past_8_throughput[i] /= M_IN_K
            past_8_throughput[i] /= M_IN_K

    movie_segment_size = np.zeros(4)
    if movie_name == 'tiledSkiing':
        movie_segment_size[0] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][0] / M_IN_K / M_IN_K
        movie_segment_size[1] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][1] / M_IN_K / M_IN_K
        movie_segment_size[2] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][2] / M_IN_K / M_IN_K
        movie_segment_size[3] = movie_json['segment_size_bytes'][index][face_idx][tile_idx][3] / M_IN_K / M_IN_K
    else:
        movie_segment_size[0] = movie_json['segment_size_bytes'][index][face][0] / M_IN_K / M_IN_K
        movie_segment_size[1] = movie_json['segment_size_bytes'][index][face][1] / M_IN_K / M_IN_K
        movie_segment_size[2] = movie_json['segment_size_bytes'][index][face][2] / M_IN_K / M_IN_K
        movie_segment_size[3] = movie_json['segment_size_bytes'][index][face][3] / M_IN_K / M_IN_K

    state = np.zeros((S_INFO, S_LEN))
    state[0, 7] = fov_tmp # 上一个视频块比特率
    state[1, :] = past_8_throughput  # 当前缓冲大小
    state[2, :] = past_download_time  # KBps 过去8个视频块对应网络吞吐量
    state[3, :A_DIM] = movie_segment_size # 10 sec 过去8个视频块下载时间
    state[4, :] = buffer / BUFFER_NORM_FACTOR # 未来视频块不同质量大小
    state[5, 7] = bitrates[lastQuality] / float(np.max(bitrates))   # 剩余视频块数量

    # 随机性选择动作
    _, _, action_prob = actor.get_actor_out(convert_torch(np.reshape(state, (1, S_INFO, S_LEN))))
    # 随机策略
    # action_prob = action_prob.cpu()
    # m = Categorical(action_prob)
    # bit_rate = m.sample().detach().item()
    # 确定性策略
    action_prob = action_prob.numpy()
    bit_rate = action_prob.argmax()

    print("ABR request suceess", coord, tile_dev, action_prob, bit_rate)
    socketio.emit('ABR response', {'data': str(bit_rate)}, to=msg['id'])


@socketio.on('connect', namespace='/FoV')
# @app.route('/', methods=['GET'])
def home():
    global videoName, videoId, isLoad, sal_maps, time_array, FoV_net

    # videoName = request.args.get("videoName")
    # if videoName is None:
    #     return "Backend of the server! None parameter!"
    videoName = 'Skiing'

    sal_maps = sal_maps_list[videoId]
    time_array = time_array_list[videoId]
    print(f'videoName={videoName}, videoId={videoId}, '
          f'sal_maps\' shape={sal_maps.shape}, time_array\'s shape={time_array.shape}')
    if isLoad[videoId]:
        socketio.emit('name response', {'data': "Success! Have Loaded model parameters!"},
                      namespace='/FoV')
        return f"Success! Have Loaded model parameters!"

    if torch.cuda.is_available():
        FoV_net = FoV_net.cuda()
    if not os.path.exists(args.model_path):
        exit("model path doesn't exist")
    # modelPath = args.model_path + f"convlstm_offline_{videoName.lower()}_1_1s.pth"
    modelPath = args.model_path + f"convlstm_offline_{videoName.lower()}_1_4s.pth"
    if not os.path.exists(modelPath):
        exit(f"{modelPath} doesn't exit!")

    FoV_net.load_state_dict(torch.load(modelPath))
    print('Loading model from', modelPath)
    print("Total number of parameters in networks is {}".format(sum(x.numel() for x in FoV_net.parameters())))
    FoV_net.eval()
    isLoad[videoId] = True
    socketio.emit('name response', {'data': f"Success! Have gotten videoName={videoName}!"},
                   namespace='/FoV')
    return f"Success! Have gotten videoName={videoName}!"


@socketio.on('FoV request', namespace='/FoV')
def wsFoV(msg):
    # print(msg)
    start_t = time.time()
    req = msg['data']
    req.reverse()
    # for item in req:
    #     print(item)
    if not req or len(req) <= 3:
        return "the length of post_data < 4"
    inputs, pre_time = get_sal_fix(time_array, sal_maps, req)

    # return jsonify([pre_time, [-0.7,0,-0.7]])
    # output_list 每个时间点最后一层隐层结果: list
    # output_last 最后一个时间点隐层和细胞层结果: list
    # conv_output 最后一个时间点隐层和细胞层卷积之后结果: list
    # predict_array 每个时间点隐层卷积之后结果 torch.tensor
    output_list, output_last, conv_output, conv_output_list_ret = FoV_net(inputs)
    pre_fov, pre_tiles = return_fov(conv_output)

    end_t = time.time()
    print(f"Time={round(end_t - start_t, 3)}")
    ret = [pre_time, pre_fov]
    # return jsonify([pre_time, pre_fov])

    socketio.emit('FoV response', {'data': ret}, to=msg['id'], namespace='/FoV')
    print("FoV request from {} success".format(msg['id']))


if __name__ == '__main__':
    socketio.run(app, host='0.0.0.0', port=9009)
